package com.kwdb.sdk.grpc;

import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import com.kwdb.sdk.query.QueryResult;
import com.kwdb.grpc.*;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * KWDB gRPC客户端 - 通信层实现
 * 
 * 功能特性：
 * - 长连接管理
 * - 多路复用
 * - 流式传输
 * - 自动重连
 * - 压缩传输
 */
public class KWDBGrpcClient {
    
    private final ClientConfig config;
    private ManagedChannel channel;
    private KWDBServiceGrpc.KWDBServiceBlockingStub blockingStub;
    private KWDBServiceGrpc.KWDBServiceStub asyncStub;
    private final AtomicBoolean connected = new AtomicBoolean(false);
    
    public KWDBGrpcClient(ClientConfig config) {
        this.config = config;
    }
    
    /**
     * 连接到gRPC服务器
     * @throws Exception 连接异常
     */
    public void connect() throws Exception {
        try {
            // 构建gRPC通道
            ManagedChannelBuilder<?> builder = ManagedChannelBuilder
                    .forAddress(config.getHost(), config.getPort())
                    .maxInboundMessageSize(config.getMaxInboundMessageSize())
                    .maxInboundMetadataSize(config.getMaxInboundMetadataSize());
            
            // 配置Keep-Alive
            if (config.isKeepAlive()) {
                builder.keepAliveTime(config.getKeepAliveTime(), TimeUnit.SECONDS)
                       .keepAliveTimeout(config.getKeepAliveTimeout(), TimeUnit.SECONDS)
                       .keepAliveWithoutCalls(config.isKeepAliveWithoutCalls());
            }
            
            // 配置TLS
            if (config.isEnableTls()) {
                builder.useTransportSecurity();
            } else {
                builder.usePlaintext();
            }
            
            this.channel = builder.build();
            this.blockingStub = KWDBServiceGrpc.newBlockingStub(channel);
            this.asyncStub = KWDBServiceGrpc.newStub(channel);
            
            // 健康检查
            if (healthCheck()) {
                connected.set(true);
            } else {
                throw new Exception("健康检查失败");
            }
            
        } catch (Exception e) {
            connected.set(false);
            throw new Exception("gRPC连接失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 断开连接
     */
    public void disconnect() {
        if (channel != null && !channel.isShutdown()) {
            try {
                channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        connected.set(false);
    }
    
    /**
     * 批量写入数据
     * @param dataPoints 数据点列表
     * @throws Exception 写入异常
     */
    public void writeBatch(List<DataPoint> dataPoints) throws Exception {
        if (!connected.get()) {
            throw new Exception("客户端未连接");
        }
        
        try {
            WriteBatchRequest.Builder requestBuilder = WriteBatchRequest.newBuilder()
                    .setClientId(config.getClientId())
                    .setBatchSize(dataPoints.size())
                    .setTimeoutMs(config.getBatchTimeout())
                    .setCompress(config.isEnableCompression());
            
            // 添加数据点
            for (DataPoint dataPoint : dataPoints) {
                com.kwdb.grpc.DataPoint.Builder pointBuilder = com.kwdb.grpc.DataPoint.newBuilder()
                        .setMeasurement(dataPoint.getMeasurement())
                        .setTimestamp(dataPoint.getTimestamp());
                
                // 添加标签
                for (Map.Entry<String, String> tag : dataPoint.getTags().entrySet()) {
                    pointBuilder.putTags(tag.getKey(), tag.getValue());
                }
                
                // 添加字段
                for (Map.Entry<String, Object> field : dataPoint.getFields().entrySet()) {
                    pointBuilder.putFields(field.getKey(), field.getValue().toString());
                }
                
                requestBuilder.addPoints(pointBuilder.build());
            }
            
            WriteBatchRequest request = requestBuilder.build();
            WriteResponse response = blockingStub.writeBatch(request);
            
            if (!response.getSuccess()) {
                throw new Exception("批量写入失败: " + response.getMessage());
            }
            
        } catch (StatusRuntimeException e) {
            throw new Exception("gRPC调用失败: " + e.getStatus().getDescription(), e);
        }
    }
    
    /**
     * 异步批量写入数据
     * @param dataPoints 数据点列表
     * @return CompletableFuture
     */
    public CompletableFuture<Void> writeBatchAsync(List<DataPoint> dataPoints) {
        return CompletableFuture.runAsync(() -> {
            try {
                writeBatch(dataPoints);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }
    
    /**
     * 流式写入数据
     * @param dataPoints 数据点列表
     * @throws Exception 写入异常
     */
    public void writeStream(List<DataPoint> dataPoints) throws Exception {
        if (!connected.get()) {
            throw new Exception("客户端未连接");
        }
        
        CompletableFuture<WriteResponse> future = new CompletableFuture<>();
        
        StreamObserver<WriteRequest> requestObserver = asyncStub.writeStream(
                new StreamObserver<WriteResponse>() {
                    @Override
                    public void onNext(WriteResponse response) {
                        if (!response.getSuccess()) {
                            future.completeExceptionally(
                                    new Exception("流式写入失败: " + response.getMessage()));
                        } else {
                            future.complete(response);
                        }
                    }
                    
                    @Override
                    public void onError(Throwable t) {
                        future.completeExceptionally(new Exception("流式写入错误: " + t.getMessage(), t));
                    }
                    
                    @Override
                    public void onCompleted() {
                        // 流式写入完成
                    }
                });
        
        // 分批发送数据
        int batchSize = Math.min(config.getBatchSize(), 1000);
        for (int i = 0; i < dataPoints.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, dataPoints.size());
            List<DataPoint> batch = dataPoints.subList(i, endIndex);
            
            WriteRequest.Builder requestBuilder = WriteRequest.newBuilder()
                    .setClientId(config.getClientId())
                    .setCompress(config.isEnableCompression());
            
            for (DataPoint dataPoint : batch) {
                com.kwdb.grpc.DataPoint.Builder pointBuilder = com.kwdb.grpc.DataPoint.newBuilder()
                        .setMeasurement(dataPoint.getMeasurement())
                        .setTimestamp(dataPoint.getTimestamp());
                
                // 添加标签和字段
                for (Map.Entry<String, String> tag : dataPoint.getTags().entrySet()) {
                    pointBuilder.putTags(tag.getKey(), tag.getValue());
                }
                
                for (Map.Entry<String, Object> field : dataPoint.getFields().entrySet()) {
                    pointBuilder.putFields(field.getKey(), field.getValue().toString());
                }
                
                requestBuilder.addPoints(pointBuilder.build());
            }
            
            requestObserver.onNext(requestBuilder.build());
        }
        
        requestObserver.onCompleted();
        
        // 等待完成
        try {
            future.get(config.getWriteTimeout(), TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new Exception("流式写入超时或失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 执行SQL查询
     * @param sql SQL语句
     * @return 查询结果
     * @throws Exception 查询异常
     */
    public QueryResult query(String sql) throws Exception {
        if (!connected.get()) {
            throw new Exception("客户端未连接");
        }
        
        try {
            QueryRequest request = QueryRequest.newBuilder()
                    .setSql(sql)
                    .setClientId(config.getClientId())
                    .setLimit(10000)
                    .build();
            
            QueryResponse response = blockingStub.query(request);
            
            if (!response.getSuccess()) {
                throw new Exception("查询失败: " + response.getMessage());
            }
            
            return QueryResult.fromGrpcResponse(response);
            
        } catch (StatusRuntimeException e) {
            throw new Exception("gRPC查询失败: " + e.getStatus().getDescription(), e);
        }
    }
    
    /**
     * 异步执行SQL查询
     * @param sql SQL语句
     * @return CompletableFuture<QueryResult>
     */
    public CompletableFuture<QueryResult> queryAsync(String sql) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                return query(sql);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }
    
    /**
     * 流式查询
     * @param sql SQL语句
     * @return 查询结果
     * @throws Exception 查询异常
     */
    public QueryResult queryStream(String sql) throws Exception {
        if (!connected.get()) {
            throw new Exception("客户端未连接");
        }
        
        CompletableFuture<QueryResult> future = new CompletableFuture<>();
        List<String> columns = new ArrayList<>();
        List<QueryResult.Row> allRows = new ArrayList<>();
        
        StreamObserver<QueryResponse> responseObserver = new StreamObserver<QueryResponse>() {
            @Override
            public void onNext(QueryResponse response) {
                if (response.getSuccess()) {
                    if (columns.isEmpty() && !response.getColumnsList().isEmpty()) {
                        columns.addAll(response.getColumnsList());
                    }
                    for (com.kwdb.grpc.QueryRow grpcRow : response.getRowsList()) {
                        allRows.add(QueryResult.Row.fromGrpcRow(grpcRow, columns));
                    }
                } else {
                    future.completeExceptionally(
                            new Exception("流式查询失败: " + response.getMessage()));
                }
            }
            
            @Override
            public void onError(Throwable t) {
                future.completeExceptionally(new Exception("流式查询错误: " + t.getMessage(), t));
            }
            
            @Override
            public void onCompleted() {
                QueryResult result = new QueryResult(columns, allRows, allRows.size(), 0, false);
                future.complete(result);
            }
        };
        
        QueryRequest request = QueryRequest.newBuilder()
                .setSql(sql)
                .setClientId(config.getClientId())
                .setStream(true)
                .build();
        
        asyncStub.queryStream(request, responseObserver);
        
        try {
            return future.get(config.getReadTimeout(), TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new Exception("流式查询超时或失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 健康检查
     * @return 是否健康
     */
    public boolean healthCheck() {
        try {
            HealthRequest request = HealthRequest.newBuilder()
                    .setService("KWDBService")
                    .build();
            
            HealthResponse response = blockingStub.healthCheck(request);
            return response.getStatus() == HealthResponse.ServingStatus.SERVING;
            
        } catch (Exception e) {
            return false;
        }
    }
    
    /**
     * 获取统计信息
     * @return 统计信息
     */
    public String getStats() {
        try {
            StatsRequest request = StatsRequest.newBuilder()
                    .setClientId(config.getClientId())
                    .build();
            
            StatsResponse response = blockingStub.getStats(request);
            return String.format("请求数: %d, 写入数: %d, 查询数: %d, 连接数: %d",
                    response.getTotalRequests(),
                    response.getTotalWrites(),
                    response.getTotalQueries(),
                    response.getActiveConnections());
            
        } catch (Exception e) {
            return "获取统计信息失败: " + e.getMessage();
        }
    }
    
    /**
     * 检查连接状态
     * @return 是否已连接
     */
    public boolean isConnected() {
        return connected.get() && channel != null && !channel.isShutdown();
    }
}
